#!/bin/bash


# 环境: /apdcephfs_qy3/share_976139/users/kaixunhuang/tools/env/megatron_omni_train
# 定义变量
GPUS_PER_NODE=8  # 每个节点上的 GPU 数量
MASTER_PORT=12733  # 指定端口
# DATA_PATH=/apdcephfs/private_kaixunhuang/data/chat/train_set/omni_train_2025_7-17/test/train_test.jsonl
DATA_PATH=/mnt/apdcephfs_sgfd/share_303841515/Tealab/user/xuelonggeng/data/wenetspeech_data/data_M.jsonl
VALID_DATA_PATH=/mnt/apdcephfs_sgfd/share_303841515/Tealab/user/xuelonggeng/data/wenetspeech_data/data_dev_tmp.jsonl
DIR=/mnt/apdcephfs_sgfd/share_303841515/Tealab/user/xuelonggeng/ckpt/qwen2_5omni_thinker
SAVE_PREFIX=qwen2_5omni_thinker_train_v0.31_xlgeng
SCRIPT="scripts/train_qwen2_5omni_thinker.sh"
HOSTS_WITH_PASSWORD_FILE="./hosts_with_password"  # 新的hosts文件
EXPECT_SCRIPT="remote_exec.exp"  # expect脚本路径

. utils/parse_options.sh

# 检查expect脚本是否存在
if [ ! -f "$EXPECT_SCRIPT" ]; then
    echo "错误：expect脚本 $EXPECT_SCRIPT 不存在"
    exit 1
fi

# 解析hosts_with_password文件，格式为: ip user password [slots]
declare -A NODES_INFO  # 存储节点信息的关联数组
while read -r line; do
    [[ -z "$line" ]] && continue  # 跳过空行
    read -r ip user password slots <<< "$line"
    slots=${slots:-$GPUS_PER_NODE}  # 如果没有指定slots，使用默认值
    NODES_INFO["$ip"]="$user $password $slots"
done < "$HOSTS_WITH_PASSWORD_FILE"
NODES_INFO["28.48.49.61"]="root wDK4EFy6lkBPFZ8, 测试机器"
# 生成hostfile
{
    for ip in "${!NODES_INFO[@]}"; do
        slots=$(echo "${NODES_INFO[$ip]}" | awk '{print $3}')
        echo "$ip slots=$slots"
    done
} > "${SAVE_PREFIX}.hostfile"

HOSTFILE="${SAVE_PREFIX}.hostfile"

# 解析hostfile
NODES=($(awk '{print $1}' "$HOSTFILE"))
NNODES=${#NODES[@]}
SLOTS=($(awk '{print $2}' "$HOSTFILE" | sed 's/slots=//g'))
MASTER_ADDR=${NODES[0]}

trap "kill_slaves" SIGINT

kill_slaves() {
    echo "捕获到 Ctrl+C，终止从节点上的训练进程..."
    for ip in "${!NODES_INFO[@]}"; do
        read -r user password _ <<< "${NODES_INFO[$ip]}"
        echo "正在终止 $ip 上的进程..."
        /usr/bin/expect "$EXPECT_SCRIPT" "$ip" "$user" "$password" \
            "ps -ef | grep finetune_audiollm | awk '{print \$2}' | xargs kill -9" &
    done
    wait
    echo "所有从节点进程已终止"
    exit 1
}

# 创建日志目录
mkdir -p logs

# 启动训练进程
for i in "${!NODES[@]}"; do
    NODE_RANK=$i
    NODE_IP=${NODES[$i]}
    # GPUS_PER_NODE=${SLOTS[$i]}
    WORLD_SIZE=$(($GPUS_PER_NODE * $NNODES))
    
    read -r user password _ <<< "${NODES_INFO[$NODE_IP]}"
    echo "使用用户 $user 登录节点 $NODE_IP"
    
    CMD="cd $PWD; bash -x $SCRIPT --DIR $DIR --MASTER_ADDR $MASTER_ADDR --MASTER_PORT $MASTER_PORT --NNODES $NNODES --NODE_RANK $NODE_RANK --WORLD_SIZE $WORLD_SIZE --DATA_PATH $DATA_PATH --VALID_DATA_PATH $VALID_DATA_PATH --SAVE_PREFIX $SAVE_PREFIX 2>&1 | tee logs/${SAVE_PREFIX}.$NODE_RANK.log"
    # CMD="cd $PWD; hostname > logs/test_$NODE_RANK.log"
    
    echo "在节点 $NODE_IP (rank $NODE_RANK) 上启动训练..."
    /usr/bin/expect "$EXPECT_SCRIPT" "$NODE_IP" "$user" "$password" "$CMD" &
done

wait
echo "所有节点上的训练已启动"
